from pyspark.sql import SparkSession
from pyspark.sql.functions import col, avg, max, min, count, round
import matplotlib.pyplot as plt
import pandas as pd
import os
from matplotlib import font_manager
import seaborn as sns
import numpy as np

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # macOS的通用中文字体
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 设置数据目录的绝对路径
DATA_DIR = "/Users/guopeiran/Code/VSCode/Something/something/stockdata"

# 创建SparkSession
spark = SparkSession.builder \
    .appName("Stock Analysis") \
    .getOrCreate()

# 读取所有股票数据
stock_data = spark.read.csv(os.path.join(DATA_DIR, "*.csv"), header=True, inferSchema=True)

def plot_simple_price_trend(stock_code):
    """绘制简化的价格走势图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    plt.plot(stock_pdf['date'], stock_pdf['close'], label='收盘价', linewidth=2, color='#2ecc71')
    
    # 添加趋势说明
    first_price = stock_pdf['close'].iloc[0]
    last_price = stock_pdf['close'].iloc[-1]
    change_percent = ((last_price - first_price) / first_price) * 100
    
    trend_text = f"整体趋势: {'上涨' if change_percent > 0 else '下跌'} {abs(change_percent):.1f}%"
    plt.text(0.02, 0.98, trend_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票价格走势图', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('价格', fontsize=12)
    
    # 优化日期标签显示
    ax = plt.gca()
    # 获取当前的日期标签位置
    xticks = ax.get_xticks()
    # 选择合适数量的标签（比如10个）
    n_labels = 10
    step = max(len(stock_pdf) // n_labels, 1)
    # 设置新的标签位置和标签
    ax.set_xticks(range(0, len(stock_pdf), step))
    ax.set_xticklabels(stock_pdf['date'].iloc[::step], rotation=45, ha='right')
    
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_price_trend.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_volume(stock_code):
    """绘制简化的成交量图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    # 根据涨跌设置颜色
    colors = ['#e74c3c' if x > 0 else '#2ecc71' for x in stock_pdf['p_change']]
    plt.bar(range(len(stock_pdf)), stock_pdf['volume'], alpha=0.6, color=colors)
    
    # 添加成交量说明
    avg_volume = stock_pdf['volume'].mean()
    volume_text = f"平均成交量: {avg_volume:,.0f}"
    plt.text(0.02, 0.98, volume_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票成交量分析\n(红色表示上涨，绿色表示下跌)', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('成交量', fontsize=12)
    
    # 优化日期标签显示
    ax = plt.gca()
    # 选择合适数量的标签（比如10个）
    n_labels = 10
    step = max(len(stock_pdf) // n_labels, 1)
    # 设置新的标签位置和标签
    ax.set_xticks(range(0, len(stock_pdf), step))
    ax.set_xticklabels(stock_pdf['date'].iloc[::step], rotation=45, ha='right')
    
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_volume_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_price_distribution(stock_code):
    """绘制简化的价格分布图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    plt.figure(figsize=(10, 6))
    sns.histplot(data=stock_pdf, x='close', bins=30, color='#3498db')
    
    # 添加价格说明
    price_text = f"平均价格: {stock_pdf['close'].mean():.2f}\n最高价: {stock_pdf['close'].max():.2f}\n最低价: {stock_pdf['close'].min():.2f}"
    plt.text(0.02, 0.98, price_text, transform=plt.gca().transAxes, 
             fontsize=12, verticalalignment='top')
    
    plt.title(f'{stock_code}股票收盘价分布', fontsize=14, pad=15)
    plt.xlabel('价格', fontsize=12)
    plt.ylabel('出现次数', fontsize=12)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_price_distribution.png', dpi=300, bbox_inches='tight')
    plt.close()

def plot_simple_trend(stock_code):
    """绘制简化的趋势图"""
    stock_df = spark.read.csv(os.path.join(DATA_DIR, f"{stock_code}.csv"), header=True, inferSchema=True)
    stock_pdf = stock_df.toPandas()
    
    # 计算5日和20日均线
    stock_pdf['MA5'] = stock_pdf['close'].rolling(window=5).mean()
    stock_pdf['MA20'] = stock_pdf['close'].rolling(window=20).mean()
    
    plt.figure(figsize=(15, 6))  # 增加图表宽度
    plt.plot(range(len(stock_pdf)), stock_pdf['close'], label='收盘价', linewidth=2, color='#2ecc71')
    plt.plot(range(len(stock_pdf)), stock_pdf['MA5'], label='5日均线', linewidth=1, color='#e74c3c')
    plt.plot(range(len(stock_pdf)), stock_pdf['MA20'], label='20日均线', linewidth=1, color='#3498db')
    
    # 添加趋势说明
    trend_text = "趋势说明:\n- 绿线：每日收盘价\n- 红线：5日均线（短期趋势）\n- 蓝线：20日均线（长期趋势）"
    plt.text(0.02, 0.98, trend_text, transform=plt.gca().transAxes, 
             fontsize=10, verticalalignment='top')
    
    plt.title(f'{stock_code}股票趋势分析', fontsize=14, pad=15)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('价格', fontsize=12)
    
    # 优化日期标签显示
    ax = plt.gca()
    # 选择合适数量的标签（比如10个）
    n_labels = 10
    step = max(len(stock_pdf) // n_labels, 1)
    # 设置新的标签位置和标签
    ax.set_xticks(range(0, len(stock_pdf), step))
    ax.set_xticklabels(stock_pdf['date'].iloc[::step], rotation=45, ha='right')
    
    plt.legend(fontsize=10)
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f'{stock_code}_trend_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

def analyze_single_stock(stock_code):
    """分析单支股票并生成图表"""
    # 生成各种分析图表
    plot_simple_price_trend(stock_code)
    plot_simple_volume(stock_code)
    plot_simple_price_distribution(stock_code)
    plot_simple_trend(stock_code)

def analyze_all_stocks():
    """分析所有股票并找出最佳投资日期"""
    # 计算每支股票的波动率
    stock_volatility = stock_data.groupBy("date").agg(
        ((max("high") - min("low")) / avg("close") * 100).alias("volatility")
    )
    
    # 找出波动率最小的股票
    best_stock = stock_volatility.orderBy("volatility").first()
    
    return best_stock

def plot_simple_summary(best_stock_date):
    """生成简化的综合分析报告"""
    # 获取推荐日期的所有股票数据
    recommended_stocks = stock_data.filter(col("date") == best_stock_date)
    recommended_pdf = recommended_stocks.toPandas()
    
    # 创建子图布局
    fig = plt.figure(figsize=(15, 10))
    gs = fig.add_gridspec(2, 2)
    
    # 1. 价格分布图
    ax1 = fig.add_subplot(gs[0, 0])
    sns.histplot(data=recommended_pdf, x='close', bins=30, color='#3498db', ax=ax1)
    ax1.set_title('推荐日期股票价格分布', fontsize=12, pad=10)
    ax1.set_xlabel('价格')
    ax1.set_ylabel('股票数量')
    ax1.grid(True, linestyle='--', alpha=0.7)
    
    # 2. 涨跌分布图
    ax2 = fig.add_subplot(gs[0, 1])
    colors = ['#e74c3c' if x > 0 else '#2ecc71' for x in recommended_pdf['p_change']]
    ax2.bar(range(len(recommended_pdf)), recommended_pdf['p_change'], color=colors)
    ax2.set_title('推荐日期股票涨跌分布\n(红色表示上涨，绿色表示下跌)', fontsize=12, pad=10)
    ax2.set_xlabel('股票序号')
    ax2.set_ylabel('涨跌幅(%)')
    ax2.grid(True, linestyle='--', alpha=0.7)
    
    # 3. 成交量分布图
    ax3 = fig.add_subplot(gs[1, 0])
    sns.histplot(data=recommended_pdf, x='volume', bins=30, color='#9b59b6', ax=ax3)
    ax3.set_title('推荐日期股票成交量分布', fontsize=12, pad=10)
    ax3.set_xlabel('成交量')
    ax3.set_ylabel('股票数量')
    ax3.grid(True, linestyle='--', alpha=0.7)
    
    # 4. 关键指标统计
    ax4 = fig.add_subplot(gs[1, 1])
    ax4.axis('off')
    stats_text = f"""
    推荐日期: {best_stock_date}
    
    市场概况:
    - 股票数量: {len(recommended_pdf)}支
    - 平均价格: {recommended_pdf['close'].mean():.2f}元
    - 最高价格: {recommended_pdf['high'].max():.2f}元
    - 最低价格: {recommended_pdf['low'].min():.2f}元
    
    涨跌情况:
    - 平均涨跌幅: {recommended_pdf['p_change'].mean():.2f}%
    - 上涨股票: {len(recommended_pdf[recommended_pdf['p_change'] > 0])}支
    - 下跌股票: {len(recommended_pdf[recommended_pdf['p_change'] < 0])}支
    
    交易情况:
    - 平均成交量: {recommended_pdf['volume'].mean():,.0f}手
    """
    ax4.text(0.1, 0.5, stats_text, fontsize=12, va='center')
    
    # 设置总标题
    fig.suptitle('推荐股票市场分析报告', fontsize=16, y=0.95)
    
    # 调整布局
    plt.tight_layout()
    plt.savefig('recommended_stocks_analysis.png', dpi=300, bbox_inches='tight')
    plt.close()

# 主函数
def main():
    # 分析单支股票（以600415为例）
    analyze_single_stock("600415")
    
    # 分析所有股票
    best_stock = analyze_all_stocks()
    
    # 生成推荐股票的综合分析图表
    plot_simple_summary(best_stock['date'])
    
    print("\n推荐投资股票分析：")
    print(f"根据分析，推荐投资日期为 {best_stock['date']} 的股票，")
    print(f"该日期的股票波动率为 {best_stock['volatility']:.2f}%，")
    print("表明该时期市场相对稳定，投资风险较低。")
    
    print("\n已生成以下分析图表：")
    print("1. 600415_price_trend.png - 价格走势图（带趋势说明）")
    print("2. 600415_volume_analysis.png - 成交量分析图（红涨绿跌）")
    print("3. 600415_price_distribution.png - 价格分布图（带统计说明）")
    print("4. 600415_trend_analysis.png - 趋势分析图（带均线说明）")
    print("5. recommended_stocks_analysis.png - 市场分析报告")

if __name__ == "__main__":
    main() 